/*
# ehh.cpp: the Ehh shared object source file.
#
# Copyright (C) 2014 Daniel Gutson, Taller Technologies
#
# This file is part of the Ehh project - the exec hook hacker.
#
# Ehh is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ehh is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with fudepan-build.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <dlfcn.h>
#include <stdio.h>  // for printf
#include <string.h> // for strdup
#include <string>
#include <sstream>
#include <fstream>
#include <list>
#include <unordered_map>
#include <forward_list>
#include <algorithm>

using namespace std;

#ifdef DEBUG_EHH
#   define execve   execve_debug
#   define execv    execv_debug
#   define execp    execvp_debug
static void show(const char* func, const char* file, char* const argv[]);
#else
#   define show(func, file, argv)
#endif

extern "C" {

#ifdef __THROW
#   define EHH_GLIBC_NOTHROW   __THROW
#else
#   define EHH_GLIBC_NOTHROW
#endif

int execve (const char* file, char *const argv[], char* const envp[]) EHH_GLIBC_NOTHROW;
int execv(const char* path, char* const argv[]) EHH_GLIBC_NOTHROW;
int execvp(const char* file, char* const argv[]) EHH_GLIBC_NOTHROW;

typedef int ExecveType (const char* file, char *const argv[], char* const envp[]);
typedef int ExecvType(const char* path, char* const argv[]);
typedef int ExecvpType(const char* file, char* const argv[]);

}

static const char* LIBC_CFG = HOME "/.ehh/libc.cfg";
static const char* EHH_RC = HOME "/.ehh/.ehh_rc";
//static const char* LIBC_CFG = "~/.ehh/libc.cfg";
//static const char* EHH_RC = "~/.ehh/.ehh_rc";

class Ehh
{
    typedef string LibcName;
    typedef string BinName;

    class DlInfo
    {
    public:
        void* dl_handle;
        void* execve_orig;
        void* execv_orig;
        void* execvp_orig;

        DlInfo() :
            dl_handle(nullptr),
            execve_orig(nullptr),
            execv_orig(nullptr),
            execvp_orig(nullptr)
        {}

        ~DlInfo()
        {
            if (dl_handle != nullptr)
                dlclose(dl_handle);
        }

        bool load(const char* dlname)
        {
            dl_handle = dlopen(dlname, RTLD_LAZY);
            bool ret = (dl_handle != nullptr);

            if (ret)
            {
                ret = loadSym(dlname, "execve", execve_orig) &&
                      loadSym(dlname, "execv", execv_orig) &&
                      loadSym(dlname, "execvp", execvp_orig);
            }
            else
                fprintf(stderr, "Ehh: Error '%s' loading %s\n", dlerror(), dlname);

            return ret;
        }
    private:
        bool loadSym(const char* dlname, const char* symname, void*& ptr) const
        {
            ptr = dlsym(dl_handle, symname);
            const bool ret = (ptr != nullptr);
            if (!ret)
                fprintf(stderr, "Ehh: error '%s' while loading %s from %s\n", dlerror(), symname, dlname);
            return ret;
        }
    };

    struct TranslationRule
    {
        virtual void apply(list<string>& args) = 0;
        virtual ~TranslationRule(){}
    };

    class IfExistsReplaceRule : public TranslationRule
    {
        const string old_arg;
        const string new_arg;
        virtual void apply(list<string>& args)
        {
            replace(args.begin(), args.end(), old_arg, new_arg);
        }
    public:
        IfExistsReplaceRule(const string& old_arg, const string& new_arg) :
            old_arg(old_arg), new_arg(new_arg)
        {}
    };

    class IfExistsRemoveRule : public TranslationRule
    {
        const string target_arg;
        virtual void apply(list<string>& args)
        {
            args.remove(target_arg);
        }
    public:
        IfExistsRemoveRule(const string& target_arg) :
            target_arg(target_arg)
        {}
    };

    class IfNotExistsAddRule : public TranslationRule
    {
        const string target_arg;
        virtual void apply(list<string>& args)
        {
            if (find(args.begin(), args.end(), target_arg) == args.end())
                args.push_back(target_arg);
        }
    public:
        IfNotExistsAddRule(const string& target_arg) :
            target_arg(target_arg)
        {}
    };

    class TranslationRules
    {
        forward_list<TranslationRule*> rules;
        static bool getOneArg(stringstream& ss, string& arg)
        {
            return (ss >> arg) && ss.str().empty();
        }

        static bool getTwoArgs(stringstream& ss, string& arg1, string& arg2)
        {
            return (ss >> arg1 >> arg2) && ss.str().empty();
        }

        template <class RuleType>
        bool addOneArgRuleType(stringstream& ss)
        {
            string arg;
            const bool ret = (ss >> arg) && ss.eof();
            if (ret)
                rules.push_front(new RuleType(arg));
            return ret;
        }

        template <class RuleType>
        bool addTwoArgsRuleType(stringstream& ss)
        {
            string arg1;
            string arg2;
            const bool ret = (ss >> arg1 >> arg2) && ss.eof();
            if (ret)
                rules.push_front(new RuleType(arg1, arg2));
            return ret;
        }
    public:
        bool addRuleLine(stringstream& ss)
        {
            string ruleType;
            bool ret = (ss >> ruleType);
            if (ret)
            {
                if (ruleType == "IF_EXISTS_REPLACE")
                    ret = addTwoArgsRuleType<IfExistsReplaceRule>(ss);
                else if (ruleType == "IF_EXISTS_REMOVE")
                    ret = addOneArgRuleType<IfExistsRemoveRule>(ss);
                else if (ruleType == "IF_NOT_EXISTS_ADD")
                    ret = addOneArgRuleType<IfNotExistsAddRule>(ss);
            }
            return ret;
        }

        void apply(char *const argv[], char**& argv2) const
        {
            list<string> args;
            decltype(args.size()) i = 0u;

            while (argv[i] != nullptr)
            {
                args.push_back(argv[i]);
                ++i;
            }

            for (auto it : rules)
                it->apply(args);

            argv2 = new char*[args.size() + 1u];    // +1 for the null
            i = 0u;
            for (auto& it : args)
                argv2[i++] = strdup(it.c_str());

            argv2[args.size()] = nullptr;
        }

        ~TranslationRules()
        {
            for (auto it : rules)
                delete it;
        }
    };

    struct BinInfo
    {
        const DlInfo* dlinfo;
        TranslationRules rules;
    };

    unordered_map<BinName, BinInfo*> binNameInfoMapping;        // own BinInfo
    unordered_map<BinName, BinInfo*> baseBinNameInfoMapping;    // do not own BinInfo
    unordered_map<LibcName, DlInfo*> libcNameDlinfoMapping;
    const DlInfo* defaultDlInfo;  // do not delete

    void createBinInfo(const BinName& binName, DlInfo* dlinfo)
    {
        if (binName == "*")
            defaultDlInfo = dlinfo;
        else
        {
            BinInfo* const binInfo = new BinInfo;
            binInfo->dlinfo = dlinfo;
            binNameInfoMapping[binName] = binInfo;
            // add the basename too:
            const auto slash_pos = binName.rfind('/');
            if (slash_pos != string::npos)
                baseBinNameInfoMapping[binName.substr(slash_pos+1)] = binInfo;
        }
    }

    bool loadLibcCfg()
    {
        ifstream file(LIBC_CFG);
        const bool ret = file;
        if (ret)
        {
            BinName binName;
            LibcName libcName;
            while (file >> binName >> libcName)
            {
                DlInfo* dlinfo;
                const auto it = libcNameDlinfoMapping.find(libcName);
                if (it == libcNameDlinfoMapping.end())
                {
                    dlinfo = new DlInfo;
                    if (dlinfo->load(libcName.c_str()))
                    {
                        libcNameDlinfoMapping[libcName] = dlinfo;
                    }
                    else
                    {
                        delete dlinfo;  // failed.
                        dlinfo = nullptr;
                    }
                }
                else
                    dlinfo = it->second;

                if (dlinfo != nullptr)
                    createBinInfo(binName, dlinfo);
            }
        }
        else
            fprintf(stderr, "%s not found\n", LIBC_CFG);

        return ret;
    }

    BinInfo* findBinary(const BinName& binName) const
    {
        BinInfo* ret = nullptr;
        auto binInfo = binNameInfoMapping.find(binName);
        if (binInfo != binNameInfoMapping.end())
            ret = binInfo->second;
        else
        {
            binInfo = baseBinNameInfoMapping.find(binName);
            if (binInfo != baseBinNameInfoMapping.end())
                ret = binInfo->second;
        }

        return ret;
    }

    void processEhhRcLine(const string& line)
    {
        /* Format:
                bin condition_action argument1 [argument2]
        */
        stringstream ss(line);
        BinName binName;
        if (ss >> binName)
        {
            auto binInfo = findBinary(binName);
            if (binInfo != nullptr)
            {
                if (!binInfo->rules.addRuleLine(ss))
                    fprintf(stderr, "Invalid rule or syntax in line '%s' in %s.\n", line.c_str(), EHH_RC);
            }
            else
                fprintf(stderr, "Binary '%s' specified in %s not found.\n", binName.c_str(), EHH_RC);
        }
        else
            fprintf(stderr, "Malformed line '%s' in %s.\n", line.c_str(), EHH_RC);
    }

    void loadEhhRc()
    {
        ifstream file(EHH_RC);
        const bool ret = file;
        if (ret)
        {
            string line;
            while (getline(file, line))
            {
                if (!line.empty() && line[0] != '#')
                    processEhhRcLine(line);
            }
        }
        else
            fprintf(stderr, "%s not found\n", EHH_RC);
    }

    void prepareExec(const BinName& binName, char* const argv[], char**& argv2, const DlInfo*& dlinfo) const
    {
        const BinInfo* const binInfo = findBinary(binName);

        if (binInfo != nullptr)
        {
            binInfo->rules.apply(argv, argv2);
            dlinfo = binInfo->dlinfo;
        }
        else
        {
            copyArgvToArgv2(argv, argv2);
            dlinfo = defaultDlInfo;
        }
    }

    static void copyArgvToArgv2(char* const argv[], char**& argv2)
    {
        auto count = 0u;
        while (argv[count] != nullptr)
            ++count;

        argv2 = new char*[count + 1u];

        for (auto i = 0u; i < count; i++)
            argv2[i] = strdup(argv[i]);

        argv2[count] = nullptr;
    }

    static void freeArgv2(char** argv2)
    {
        auto i = 0u;
        while (argv2[i] != nullptr)
            free(argv2[i++]); // free since they where strdup-ped.

        delete [] argv2;
    }
public:
    Ehh()
    {
        if (loadLibcCfg())
            loadEhhRc();
    }

    ~Ehh()
    {
        for (auto& it : binNameInfoMapping)
            delete it.second;

        for (auto& it : libcNameDlinfoMapping)
            delete it.second;
    }

    int execve(const char* file, char* const argv[], char* const envp[])
    {
        char** argv2;
        const DlInfo* dlinfo;
        prepareExec(file, argv, argv2, dlinfo);
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-pedantic"
        const int ret = (*(reinterpret_cast<ExecveType*>(dlinfo->execve_orig)))(file, argv2, envp);
#pragma GCC diagnostic pop
        freeArgv2(argv2);
        return ret;
    }

    int execv(const char* path, char* const argv[])
    {
        char** argv2;
        const DlInfo* dlinfo;
        prepareExec(argv[0], argv, argv2, dlinfo);
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-pedantic"
        const int ret = (*(reinterpret_cast<ExecvType*>(dlinfo->execv_orig)))(path, argv2);
#pragma GCC diagnostic pop
        freeArgv2(argv2);
        return ret;
    }

    int execvp(const char* file, char* const argv[])
    {
        char** argv2;
        const DlInfo* dlinfo;
        prepareExec(file, argv, argv2, dlinfo);
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-pedantic"
        const int ret = (*(reinterpret_cast<ExecvpType*>(dlinfo->execvp_orig)))(file, argv2);
#pragma GCC diagnostic pop
        freeArgv2(argv2);
        return ret;
    }
};

static Ehh ehh;

int execve(const char *file, char* const argv[], char* const envp[]) EHH_GLIBC_NOTHROW
{
    show("execve", file, argv);
    return ehh.execve(file, argv, envp);
}

int execv(const char* path, char* const argv[]) EHH_GLIBC_NOTHROW
{
    show("execv", argv[0], argv);
    return ehh.execv(path, argv);
}

int execvp(const char* file, char* const argv[]) EHH_GLIBC_NOTHROW
{
    show("execvp", file, argv);
    return ehh.execvp(file, argv);
}

#ifdef DEBUG_EHH
void show(const char* func, const char* file, char* const argv[])
{
    fprintf(stderr, "EHH: func: %s - file %s ", func, file);
    unsigned int i = 0u;
    while (argv[i] != nullptr)
    {
        fprintf(stderr, "arg[%d]=%s ", i, argv[i]);
        ++i;
    }
    fprintf(stderr, "\n");
}

int main()
{

    char arg1[] = "/bin/ls";
    char arg2[] = "-l";
    char* const args[] = { arg1, arg2, nullptr  };
    char* const envs[] = { nullptr };

    const int ret = execve(arg1, args, envs);
    return ret;
}
#endif

